import os
from sqlalchemy.testing.plugin.plugin_base import config
from vanna.base import VannaBase
from vanna.chromadb import ChromaDB_VectorStore
from volcenginesdkarkruntime import Ark
from abc import ABC


class DouBao(VannaBase, ABC):
    def __init__(self):
        super().__init__(config)
        self.client = Ark(
        sk="AKLTOThkMDRiZWI2ZTNlNGY1YmE0NzRmNDE5OTQyYmU3ZWU",
        ak="WTJOaU5UWTFOV1V5TkRsaE5ETTNaR0ZsTURFNE9EazFNVGxsTm1JeVlqVQ==",
        api_key='78018cba-25ba-4e0b-ac6d-5ab1a2abecf6'
    )

    def assistant_message(self, message: str):
        print(f"{message}")
        return message

    def system_message(self, message: str):
        print(f"{message}")
        return message

    def user_message(self, message: str):
        print(f"{message}")
        return message


    def submit_prompt(self, prompt, **kwargs):
        if prompt is None:
            raise Exception("Prompt is None")
        if len(prompt) == 0:
            raise Exception("Prompt is empty")
        completion = self.client.chat.completions.create(
            model="ep-20240809161649-gtcl7",
            messages=[
                {"role": "system", "content": "你是豆包，是由字节跳动开发的 AI 人工智能助手"},
                {"role": "user", "content": ",".join(prompt)},
            ],
        )
        return completion.choices[0].message.content


class DouBaoVanna(ChromaDB_VectorStore, DouBao, ABC):
    def __init__(self, config_vb):
        ChromaDB_VectorStore.__init__(self, config=config_vb)
        DouBao.__init__(self)



vn =DouBaoVanna(config_vb={"path":os.path.join(os.getcwd(), "weekly_work")})
vn.connect_to_mysql(host='192.168.35.221', dbname='llm', user='root', password='xxtc.mysql2@147258', port=3307)

    # llm_sql = "CREATE TABLE  项目进度周报表 (\
    #     序号 INT PRIMARY KEY auto_increment,\
    #     研究领域 VARCHAR(255),\
    #     项目名称 VARCHAR(255),\
    #     项目负责人 VARCHAR(255),\
    #     项目类型 VARCHAR(255),\
    #     计划开始时间 DATE,\
    #     计划结束时间 DATE,\
    #     周期（天） INT,\
    #     进度核算时间 DATE,\
    #     序时进度 DECIMAL(5, 4),\
    #     实际进度 DECIMAL(3, 2),\
    #     本周工作内容 TEXT,\
    #     下周计划 TEXT,\
    #     更新时间 TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP\
    # )"
    # llm_sql = "CREATE  TABLE `寄递研究所工作安排20250311_144420_905` (\
    #   `序号` INT NOT NULL AUTO_INCREMENT,\
    #   `研究领域` VARCHAR(255),\
    #   `项目名称` VARCHAR(255),\
    #   `项目负责人` VARCHAR(255),\
    #   `项目类型` VARCHAR(255),\
    #   `计划开始时间` DATETIME,\
    #   `计划结束时间` DATETIME,\
    #   `周期（天）` INT,\
    #   `进度核算时间` DATETIME,\
    #   `序时进度（自动生成）` DECIMAL(10,4),\
    #   `实际进度` DECIMAL(10,4),\
    #   `本周工作内容` TEXT,\
    #   `下周计划` TEXT,\
    #   `更新时间` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,\
    #   PRIMARY KEY (`序号`)\
    # ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;"
    # # vn.train(ddl=llm_sql)
    #
    # # a = vn.get_related_ddl(question="规范化管理工作的项目负责人是谁")
    # # print(a)
    # # vn.ask("超市7月23日的营业收入是多少?")
    # # vn.ask("顺丰2023年的总营业收入是多少?")
    # # vn.ask(question="顺丰2022年的营业收入同比增加了多少？",allow_llm_to_see_data=True)
    # # vn.ask(question="在研有多少个项目要求为不重复的项目？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="负责两个以上项目的项目负责人有哪些人？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="规范化管理工作这个项目的每周进度情况？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="规范化管理工作这个项目的每周开展情况？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="规范化管理工作这个项目的每周工作内容？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="规范化管理工作这个项目的每周进度值？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="规范化管理工作这个项目的每周实际进度？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="规范化管理工作这个项目的每周进度？",allow_llm_to_see_data=True,visualize=False)
# vn.ask(question="规范化管理工作这个项目的每周进度？",allow_llm_to_see_data=True,visualize=False,auto_train=False)
    # # vn.ask(question="规范化管理工作这个项目的项目工作内容及项目负责人？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="负责项目最多的人是谁？",allow_llm_to_see_data=True,visualize=False)
    # # vn.ask(question="项目负责平均负责多少个项目？，注意去重",allow_llm_to_see_data=True,visualize=False)
